import json
import xml.etree.ElementTree as ET
import re
import requests
from transformers import MarianMTModel, MarianTokenizer
import os
import torch
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

model_name = "Helsinki-NLP/opus-mt-en-zh"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def local_translate(text):
    batch = tokenizer([text], return_tensors="pt")
    batch.to(device)
    gen = model.generate(**batch)
    return tokenizer.decode(gen[0], skip_special_tokens=True)

def extract_text_preserve_structure(element):
    """
    递归地从一个XML元素中提取文本内容。
    它会特别处理类似 <xhtml:br/> 的标签，将其转换成换行符，
    并尝试保留文本节点中已有的空白字符（可能代表缩进）。
    """
    parts = []
    if element.text:
        parts.append(element.text)
    
    for child in element:
        tag_local_name = child.tag
        if '}' in child.tag: # 处理命名空间，例如 {http://www.w3.org/1999/xhtml}br
            tag_local_name = child.tag.split('}', 1)[1]

        if tag_local_name.lower() == 'br':
            parts.append('\n')
        else: # 对于其他子标签 (如 xhtml:div, xhtml:i)，递归提取其内容
            parts.append(extract_text_preserve_structure(child))
        
        if child.tail:
            parts.append(child.tail)
            
    return "".join(parts)


def parse_cwe_xml(xml_file_path):
    """
    解析CWE XML文件，提取包含C语言Bad代码示例的缺陷信息。
    """
    results = []
    
    try:
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
    except ET.ParseError as e:
        print(f"XML解析错误 (XML parsing error): {e}")
        return results
    except FileNotFoundError:
        print(f"错误：文件未找到 (Error: File not found at) '{xml_file_path}'")
        return results
    # 获取命名空间
    ns = {'cwe': 'http://cwe.mitre.org/cwe-7'} if '}' in root.tag else {}
        
    # 使用命名空间查找Weaknesses
    weaknesses_element = root.find('cwe:Weaknesses', ns) if ns else root.find('Weaknesses')
    if weaknesses_element is None:
        print("错误：在XML中未找到 <Weaknesses> 标签。(Error: Could not find <Weaknesses> tag in XML.)")
        return results
    # for child in weaknesses_element:
    #     print(f"Child tag: {child.tag}, attributes: {child.attrib}")

    for weakness_node in weaknesses_element.findall('cwe:Weakness', ns):
        weakness_name = weakness_node.get('Name')
        
        description_node = weakness_node.find('cwe:Description', ns)
        description = description_node.text.strip() if description_node is not None and description_node.text else "无描述 (No description available)"
        
        bad_c_examples = []
        
        demonstrative_examples_node = weakness_node.find('cwe:Demonstrative_Examples', ns)
        if demonstrative_examples_node is not None:
            for demo_example_node in demonstrative_examples_node.findall('cwe:Demonstrative_Example', ns):
                tmp_res = {}
                intro_elem = demo_example_node.find("cwe:Intro_Text", ns)
                if intro_elem is not None and intro_elem.text:
                    intro_text = intro_elem.text.strip()
                    tmp_res["example_description"] = intro_text
                example_code_nodes = demo_example_node.findall('cwe:Example_Code', ns)
                example_codes = []
                for ec_node in example_code_nodes:
                    if ec_node.get('Language') == 'C' and ec_node.get('Nature') == 'Bad':
                        # 提取原始代码文本
                        raw_extracted_code = extract_text_preserve_structure(ec_node)
                        
                        # 清理提取的代码文本
                        # 1. 去除整个代码块前后的空白
                        code_text = raw_extracted_code.strip()
                        # 2. 将多个连续的空行（可能由多个<br/>或空白符造成）合并为一个空行
                        code_text = re.sub(r'\n\s*\n+', '\n\n', code_text) 
                        
                        if code_text: # 如果代码非空则添加
                            example_codes.append(code_text)
                if example_codes:
                    tmp_res["代码"] = example_codes
                    bad_c_examples.append(tmp_res)
                            
        if bad_c_examples:
            for item in bad_c_examples:
                if "example_description" in item:
                    item["示例描述"] = translate_text(item["example_description"])

            results.append({
                'vulnerability_name': weakness_name,
                '缺陷名称': translate_text(weakness_name),
                'description': description,
                '缺陷描述': translate_text(description),
                '代码示例': bad_c_examples
            })
            print(len(results))
    
    return results

def translate_text(text, source_lang="en", target_lang="zh"):
    """调用翻译API将文本翻译为中文(模拟实现)"""
    return local_translate(text)
    # return f"[译] {text}"
    try:
        # 实际使用时替换为真实翻译API调用
        url = "https://api.mymemory.translated.net/get"
        params = {"q": text, "langpair": f"{source_lang}|{target_lang}"}
        response = requests.get(url, params=params)
        return response.json()["responseData"]["translatedText"]
    except:
        return f"[译] {text}"

if __name__ == "__main__":
    xml_file_path = './cwec_v4.17.xml'
    # 例如: /Users/bytedance/Downloads/cwec_v4.17.xml
    # 或者在 Windows 上: C:\\path\\to\\your\\cwec_v4.17.xml

    extracted_data = parse_cwe_xml(xml_file_path)

    with open("data.json", "w", encoding="utf-8") as f:
        json.dump(extracted_data, f, ensure_ascii=False, indent=4)
    
    if extracted_data:
        print("\n--- 提取结果 ---")
        for item in extracted_data:
            print(f"\n缺陷名称 (Weakness Name): {item['vulnerability_name']}")
            print(f"描述 (Description): {item['description']}")
            print("C语言 Bad 示例 (Bad C Examples):")
            if item['代码示例']:
                for i, example in enumerate(item['代码示例']):
                    print(f"--- 示例 (Example) {i+1} ---")
                    print(example)
                    print("-------------------------")
            else:
                print("  (无C语言Bad示例) ((No C language Bad examples found for this item, though it was selected))")
            print("\n========================================")
    else:
        print("未能提取到任何包含C语言Bad示例的缺陷信息，或者文件处理失败。")
        print("(No weaknesses with C language Bad examples were extracted, or file processing failed.)")